from fastapi import FastAPI, HTTPException, Depends, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer, CrossEncoder
from transformers import AutoTokenizer
import torch
from typing import List, Union
import logging
import numpy as np

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 安全方案
security = HTTPBearer()

# 初始化模型
def load_models():
    try:
        # Embedding 模型
        embedding_model = SentenceTransformer("/mnt/remote/model_weights/bge-large-zh-v1.5")
        embedding_tokenizer = AutoTokenizer.from_pretrained("/mnt/remote/model_weights/bge-large-zh-v1.5")
        # Rerank 模型
        reranker_model = CrossEncoder(
            "/mnt/remote/model_weights/bge-reranker-large",
            default_activation_function=torch.nn.Sigmoid()
        )
        
        return {
            "embedding": (embedding_model, embedding_tokenizer),
            "reranker": reranker_model
        }
    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        raise

models = load_models()

app = FastAPI()

# # 鉴权验证
# async def verify_auth(credentials: HTTPAuthorizationCredentials = Depends(security)):
#     # 这里替换为你的实际API KEY验证逻辑
#     valid_api_key = "your-api-key"
#     if credentials.credentials != valid_api_key:
#         raise HTTPException(
#             status_code=status.HTTP_401_UNAUTHORIZED,
#             detail="无效的API密钥"
#         )
#     return True

# 数据模型
class EmbeddingRequest(BaseModel):
    input: Union[str, List[str]]
    model: str = "bge-large-zh-v1.5"
    encoding_format: str = "float"

class RerankRequest(BaseModel):
    question: str
    model: str = "bge-reranker-large"
    answers: List[str]

# Embedding 服务
@app.post("/v1/embeddings")
async def embeddings(
    request: EmbeddingRequest
):
    """生成文本嵌入向量"""
    try:
        # 验证模型
        if request.model != "bge-large-zh-v1.5":
            raise HTTPException(400, "不支持的模型")

        # 处理输入格式
        inputs = [request.input] if isinstance(request.input, str) else request.input
        
        # 计算token
        token_counts = [
            len(models["embedding"][1].encode(text, add_special_tokens=False))
            for text in inputs
        ]
        total_tokens = sum(token_counts)
        
        # 生成嵌入
        embeddings = models["embedding"][0].encode(inputs, convert_to_tensor=True)
        
        # 格式化输出
        # if request.encoding_format == "float":
        #     embeddings = embeddings.cpu().numpy().tolist()
        # else:
        #     raise HTTPException(400, "不支持的编码格式")
        embeddings = embeddings.cpu().numpy().tolist()
        return {
            "object": "list",
            "data": [
                {
                    "object": "embedding",
                    "embedding": emb,
                    "index": idx
                } for idx, emb in enumerate(embeddings)
            ],
            "model": request.model,
            "usage": {
                "prompt_tokens": total_tokens,
                "total_tokens": total_tokens
            }
        }
    except Exception as e:
        logger.error(f"嵌入生成错误: {e}")
        raise HTTPException(500, "服务内部错误")

# Rerank 服务
@app.post("/v1/reranker")
async def rerank(
    request: RerankRequest
):
    """对文档进行相关性排序"""
    try:
        if request.model != "bge-reranker-large":
            raise HTTPException(400, "不支持的模型")
        print(request)
        # 保持原始处理逻辑
        inputs = [[request.question, answer] for answer in request.answers]
        scores = models["reranker"].predict(inputs)
        
        # 仅做必要类型转换
        scores = [float(score) for score in scores]
        
        # 保持原始排序逻辑
        ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
        
        # 标准化响应格式
        return {
            "ranked_results": [{
                "id": idx, 
                "score": score
            } for idx, score in ranked],
            "model": "bge-reranker-large"
        }
    except Exception as e:
        logging.error(f"Rerank error: {e}")
        raise HTTPException(500, "Internal server error")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=10025)